import argparse
import gzip
import sys
import mappy
from collections import namedtuple
import time
import threading
import os
import multiprocessing as mp
import numpy as np
from deepsignal3.extract_features import _write_featurestr
from deepsignal3.utils.process_utils import str2bool
from deepsignal3.utils import bam_reader

def parse_args():
    parser = argparse.ArgumentParser("")
    parser.add_argument("--input",'-i', type=str,required=True)
    parser.add_argument("--output",'-o', type=str,required=True)
    parser.add_argument("--bam",'-b', type=str,required=True)
    parser.add_argument("--nproc", "-p", type=int,required=True)
    parser.add_argument("--reference_path",'-r', type=str,required=False)
    parser.add_argument("--number",'-n', type=int,required=False,default=-1)
    parser.add_argument("--w_batch_num", type=int,required=False,default=200)
    parser.add_argument("--w_is_dir", action="store",
                           type=str, required=False, default="no",
                           help='if using a dir to save features into multiple files')
    parser.add_argument("--best_n", type=int, default=1, required=False,
                            help="best_n arg in mappy(minimap2), default 1")
    return parser.parse_args()


def read_feature(input,feature_Q,number,batch_size=1,qsize_limit=40,time_wait=1):
    print('begin read feature', flush=True)
    fea_list=[]
    count=0
    with open(input, 'r') as f:
        for line in f:
            while feature_Q.qsize()>=qsize_limit:
                #print('Pausing bam and pod5 input due to OUTPUT queue size limit. Output_qsize=%d' %(output_Q.qsize()), flush=True)
                time.sleep(time_wait)
            words=line.strip().split('\t')
            seq=words[3]
            fea_list.append((words[0],words[1],words[2],seq,words[5],words[8],words[9]))
            #read_name,reference_name,str(ref_start),seq,pred_pos,mean_preds,pred_lebals
            count+=1
            if number==-1:
                if len(fea_list)>=batch_size:
                    feature_Q.put(fea_list)
                    fea_list=[]
            else:
                if count<number:
                    if len(fea_list)>=batch_size:
                        feature_Q.put(fea_list)
                        fea_list=[]
                else:               
                    break
    if len(fea_list)>0:
        feature_Q.put(fea_list)

def find_ref(bam_index,feature_Q,output_Q,time_wait=1,d_batch_size=1,qsize_limit=22):
    fea_list=[]
    while True:
        if feature_Q.empty():
            time.sleep(time_wait)
        read_data = feature_Q.get()
        if read_data == "kill":
            feature_Q.put("kill")
            break
        for (read_name,reference_name,ref_start,seq,pred_pos_text,mean_preds_text,pred_lebals_text) in read_data:
            for bam_read in bam_index.get_alignments(read_name):
                while output_Q.qsize()>=qsize_limit:
                    time.sleep(time_wait)
                if bam_read.is_supplementary or bam_read.is_secondary or bam_read.is_mapped==False:
                    continue
                if bam_read.mapping_quality<10:
                    continue
                if reference_name != bam_read.reference_name or int(ref_start)!=bam_read.reference_start:
                    continue
                pred_pos=np.array([int(x) for x in pred_pos_text.split(",")])
                mean_preds=np.array([float(x) for x in mean_preds_text.split(",")])
                pred_lebals=np.array([int(x) for x in pred_lebals_text.split(",")])
                strand = "-" if bam_read.is_reverse else "+"
                if bam_read.modified_bases != None:
                    ref_loc = bam_read.get_reference_positions(full_length=True)
                    for m, locs in bam_read.modified_bases_forward.items():
                        if m[0] == 'C' and m[2] == 'm':
                            i=0
                            for lc in locs:
                                pos=lc[0]
                                if pos in pred_pos:
                                    bkmer_start = pos - 2 if pos - 2 >= 0 else 0
                                    bkmer_end = pos + 3 if pos + 3 <= len(seq) else len(seq)
                                    if  ref_loc[pos] is None:
                                        i+=1
                                        continue
                                    fea_list.append('\t'.join([reference_name,str(ref_loc[pos]),strand,'.',read_name,'.',str(round(1-mean_preds[i],6)),str(mean_preds[i]),str(pred_lebals[i]),seq[bkmer_start:bkmer_end]]))
                                    if len(fea_list)>=d_batch_size:
                                        output_Q.put(fea_list)
                                        fea_list=[]
                                    i+=1
                
    if len(fea_list)>0:
        output_Q.put(fea_list)


def main():
    args=parse_args()
    #if args.input.endswith(".gz"):
    #    input_file = gzip.open(args.input, 'rt')
    #else:
    #    input_file = open(args.input, 'r')
    ref_path=''
    if args.reference_path is not None:
        ref_path = os.path.abspath(args.reference_path)
    if not os.path.exists(ref_path):
        print("--reference_path not set right!")
    
    feature_Q = mp.Queue()
    output_Q = mp.Queue()
    bam_index=bam_reader.ReadIndexedBam(args.bam)
    p = mp.Process(target=read_feature,args=(args.input,feature_Q,args.number),name='read_feature')
    p.daemon = True
    p.start()
    ex_dp=args.nproc - 2
    ex_procs = []
    for i in range(ex_dp):
        pb = mp.Process(target=find_ref, args=(bam_index,feature_Q,output_Q),
                          name="pb_reader")
            
        pb.daemon = True
        pb.start()
        ex_procs.append(pb)
    p_w = mp.Process(target=_write_featurestr, args=(args.output, output_Q, args.w_batch_num,
                                                     str2bool(args.w_is_dir)),
                     name="writer")
    p_w.daemon = True
    p_w.start()
    p.join()
    
    feature_Q.put("kill")
    for pb in ex_procs:
        pb.join()
    output_Q.put("kill")
    p_w.join()

if __name__ == '__main__':
    sys.exit(main())